from util import reduce_to_k_dim
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt


def do_kmeans(text_vectors, cluster_cnt):
    kmeans = KMeans(cluster_cnt, random_state=0)
    kmeans.fit(text_vectors)  # 训练模型
    labels = kmeans.predict(text_vectors)  # 预测分类
    return labels


def visualize_kmeans(text_vectors, labels, reduced_dimension, kmeans_res_path=None):
    reduced_vectors = reduce_to_k_dim(text_vectors, reduced_dimension)
    plt.scatter(reduced_vectors[:, 0], reduced_vectors[:, 1], c=labels, s=40, cmap='viridis')
    if kmeans_res_path:
        plt.savefig(kmeans_res_path)
    plt.show()

